# AUTHOR: DING
# -*- codeing = utf-8 -*-
# @Time: 2024/2/20 8:25
# @Author: 86139
# @Site:
# @File: 08-dataset.py
# @Software: PyCharm

import torchvision
from torch.utils.tensorboard import SummaryWriter
from PIL import Image

# 建立transform
dataset_transform = torchvision.transforms.Compose([
    torchvision.transforms.ToTensor()
])

# 下载数据集
train_set = torchvision.datasets.CIFAR10(root="./dataset", train=True, transform=dataset_transform,
                                         download=True)  # 5万张
test_set = torchvision.datasets.CIFAR10(root="./dataset", train=False, transform=dataset_transform,
                                        download=True)  # 1万张
'''
# 不加transform时，输出的时PIL
img, target = test_set[0]
img.show()  # 32*32像素
'''
writer = SummaryWriter("./logs")
for i in range(10):
    img_tensor, target = train_set[i]
    writer.add_image("train_set", img_tensor, i)

writer.close()
